#!/usr/bin/env python36
# -*- coding: utf-8 -*-
"""
Created on July, 2018

@author: Tangrizzly
"""

import datetime
import math
import pdb

import numpy as np
import torch
from torch import nn
from torch.nn import Module, Parameter
import torch.nn.functional as F


class GNN(Module):
    def __init__(self, hidden_size, step=1):
        super(GNN, self).__init__()
        # GGNN网络的层数
        self.step = step
        # 隐藏层特征维度
        self.hidden_size = hidden_size
        # 输入层特征维度，这里乘以2是因为输入是出度矩阵和入度矩阵拼接之后的邻接矩阵，所以需要乘以2
        self.input_size = hidden_size * 2
        # 这里乘以3是因为，分别对应Wz、Wr、Wo与Uz、Ur、Uo三个权重，在GNNCell函数中分为三份
        self.gate_size = 3 * hidden_size
        # 定义一个形状为（self.gate_size, self.input_size）可训练参数，输入矩阵进行线性变换时的权重矩阵，对应公式2、3、4中括号前半部分的可学习参数
        self.w_ih = Parameter(torch.Tensor(self.gate_size, self.input_size))
        # 隐藏层矩阵进行线性变换的权重矩阵，对应公式2、3、4中括号的后半部分可学习参数
        self.w_hh = Parameter(torch.Tensor(self.gate_size, self.hidden_size))
        # 定义一个形状为（self.gate_size,）可训练参数，输入层线性变换偏置项
        self.b_ih = Parameter(torch.Tensor(self.gate_size))
        # 隐藏层线性变换偏置项
        self.b_hh = Parameter(torch.Tensor(self.gate_size))
        # 入度矩阵线性变换偏置项
        self.b_iah = Parameter(torch.Tensor(self.hidden_size))
        # 出度矩阵线性变换偏置项
        self.b_oah = Parameter(torch.Tensor(self.hidden_size))
        # 入度矩阵线性变换函数
        self.linear_edge_in = nn.Linear(self.hidden_size, self.hidden_size, bias=True)
        # 出度矩阵线性变换函数
        self.linear_edge_out = nn.Linear(self.hidden_size, self.hidden_size, bias=True)
        # 下文没用到
        self.linear_edge_f = nn.Linear(self.hidden_size, self.hidden_size, bias=True)

    '''
    形如 [:,:,:]的表达式常备用于python中的维度采样(https://so.csdn.net/so/search?from=pc_blog_highlight&q=python)
    '''
    '''
    GNNCell(self, A, hidden) #将节点的出入度邻接矩阵及隐层矩阵输入GNNCell（门控图神经网络神经元，
    在这个函数中基于nn.Linear的实现了门控循环单元的基本结构，执行论文中的更新函数[公式1-5]），
    GNNCell返回值用于更新hidden矩阵
    '''

    def GNNCell(self, A, hidden):
        """ 门控图神经网络 逻辑单元
        :param A: 会话图对应的邻接矩阵
        :param hidden: 初始的隐藏层输入，可以看作是 hidden_0
        :return: 更新之后的节点信息
        """
        # print('GNNCell中A的内容。。。。。')
        # print(A[:, :, :])
        # print('GNNCell中A的入度矩阵。。。。。')
        # print(A[:, :, :A.shape[1]])
        print('GNNCell中A的shape：', A.shape)
        # 命令行断点调试
        # pdb.set_trace()
        # torch.matmul矩阵相乘（数学上的矩阵相乘）
        # 入度矩阵、出度矩阵，对应GRU的公式(1)
        input_in = torch.matmul(A[:, :, :A.shape[1]], self.linear_edge_in(hidden)) + self.b_iah
        input_out = torch.matmul(A[:, :, A.shape[1]: 2 * A.shape[1]], self.linear_edge_out(hidden)) + self.b_oah
        # 出入度矩阵在第2个维度上拼接之后的邻接矩阵
        inputs = torch.cat([input_in, input_out], 2)
        # 公式推理参考：https://gitee.com/pp-bb/SR-GNN-master/blob/master/SR-GNN%E5%85%AC%E5%BC%8F%E4%B8%8E%E7%9B%B8%E5%85%B3%E4%BB%A3%E7%A0%81%E7%AC%94%E8%AE%B0.md
        # 对输入和隐层进行线性变换
        gi = F.linear(inputs, self.w_ih, self.b_ih)
        gh = F.linear(hidden, self.w_hh, self.b_hh)
        # ======= 内容观察 =======
        # print('GNNCell中gi的内容。。。。。')
        # print(gi[:5])
        # print('gi形状：', gi.shape)  # gi形状： torch.Size([99, 6, 300])
        # ======= 内容观察 =======
        # gi.chunk(3, 2)表示将gi在第2个维度上分为3份
        i_r, i_i, i_n = gi.chunk(3, 2)
        h_r, h_i, h_n = gh.chunk(3, 2)
        # ======= 内容观察 =======
        # print('GNNCell中i_r的内容...')
        # print(i_r[:5])
        # print('i_r形状：', i_r.shape)    # i_r形状： torch.Size([99, 6, 100])
        # print('GNNCell中i_i的内容...')
        # print(i_i[:5])
        # print('i_i形状：', i_i.shape)    # i_i形状： torch.Size([99, 6, 100])
        # print('GNNCell中i_n的内容...')
        # print(i_n[:5])
        # print('i_n形状：', i_n.shape)    # i_n形状： torch.Size([99, 6, 100])
        # ======= 内容观察 =======
        # 重置门，公式(3)
        resetgate = torch.sigmoid(i_r + h_r)
        # 更新门，公式(2)
        inputgate = torch.sigmoid(i_i + h_i)
        # 公式(4)，得到节点邻居聚合后的信息
        newgate = torch.tanh(i_n + resetgate * h_n)
        # 这里将公式(5)进行一次拆分化简，但不完全相同 ==== 原始 =====
        hy = newgate + inputgate * (hidden - newgate)
        # ==== 原始 =====
        # ==== 修改 =====
        # 代码中的表达式展开 hy = newgate + inputgate * hidden - inputgate * newgate
        # 代码中的表达式提取公因式 hy = (1 - inputgate) * newgate + inputgate * hidden
        # 按照公式(5)这里应该是 hy = (1 - inputgate) * hidden + inputgate * newgate
        # 这里可以修改为正确的公式(5)，然后执行代码比较一下计算结果
        # ==== 修改 =====
        # print('GNNCell中hy的内容...')
        # print(hy[:5])
        print('GNNCell中hy的形状：', hy.shape)
        return hy

    def forward(self, A, hidden):
        # step是GRU网络的层数，循环之后得到step层GRU之后的结果
        for i in range(self.step):
            hidden = self.GNNCell(A, hidden)
        # 至此，得到会话图中项目的表示
        return hidden


class SessionGraph(Module):
    def __init__(self, opt, n_node):
        super(SessionGraph, self).__init__()
        self.hidden_size = opt.hiddenSize
        self.n_node = n_node
        self.batch_size = opt.batchSize
        # 是否只使用全局偏好来预测
        self.nonhybrid = opt.nonhybrid
        # 嵌入层
        self.embedding = nn.Embedding(self.n_node, self.hidden_size)
        # 获得GGNN模型
        self.gnn = GNN(self.hidden_size, step=opt.step)
        # 计算预测概率过程中的线性变换函数
        self.linear_one = nn.Linear(self.hidden_size, self.hidden_size, bias=True)
        self.linear_two = nn.Linear(self.hidden_size, self.hidden_size, bias=True)
        self.linear_three = nn.Linear(self.hidden_size, 1, bias=False)
        # 拼接长短期偏好的线性变换函数
        self.linear_transform = nn.Linear(self.hidden_size * 2, self.hidden_size, bias=True)
        # 损失函数
        self.loss_function = nn.CrossEntropyLoss()
        # 优化器
        self.optimizer = torch.optim.Adam(self.parameters(), lr=opt.lr, weight_decay=opt.l2)
        # 学习率调整，等间隔调整学习率，调整倍数为gamma倍，调整间隔为step_size，参考：https://zhuanlan.zhihu.com/p/69411064
        self.scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer, step_size=opt.lr_dc_step, gamma=opt.lr_dc)
        self.reset_parameters()

    def reset_parameters(self):
        """重置参数到区间内：(-1√hidden_size,1√hidden_size)
        """
        stdv = 1.0 / math.sqrt(self.hidden_size)
        for weight in self.parameters():
            weight.data.uniform_(-stdv, stdv)

    def compute_scores(self, hidden, mask):
        """
        计算评分
        :param hidden: (100,16,100) 其中16代表现样本所有数据最长会话的长度，长度不足其余部分补了0
        :param mask: (100,16) 有序列的位置是[1],没有动作序列的位置是[0]
        :return:
        """
        # ht 为最后一个点击项的向量表示，作为局部偏好，sum(mask, 1)：行求和，得到一列
        ht = hidden[torch.arange(mask.shape[0]).long(), torch.sum(mask, 1) - 1]  # batch_size x latent_size
        # print('compute_scores中ht的内容。。。。。')
        # print(ht)
        # print('ht的形状', ht.shape)

        # q1是公式6中的vn，q2是公式6中vi组成的矩阵
        q1 = self.linear_one(ht).view(ht.shape[0], 1, ht.shape[1])  # batch_size x 1 x latent_size
        q2 = self.linear_two(hidden)  # batch_size x seq_length x latent_size
        # 计算注意力分数，原文中公式(6)
        alpha = self.linear_three(torch.sigmoid(q1 + q2))
        # print('compute_scores中alpha的内容。。。。。')
        # print(alpha[:5])
        # 注意力机制进行加权求和，原文中公式(6)
        a = torch.sum(alpha * hidden * mask.view(mask.shape[0], -1, 1).float(), 1)
        print('a的大小：', a.shape)
        # 判断是否采用局部偏好和全局偏好结合，self.nonhybrid 默认为false，即采用局部偏好和全局偏好作为用户的最终偏好；为true则只采用全局偏好
        if not self.nonhybrid:
            # 原文中公式(7)
            a = self.linear_transform(torch.cat([a, ht], 1))
            print('a的大小：', a.shape)
        print('self.embedding.weight的大小：', self.embedding.weight.shape)  # 310x100
        # print('self.embedding.weight[0]的内容：', self.embedding.weight[0])  # 100x1
        # 一个batch内所有节点候选项的初始embedding，sample数据集中项目个数为310
        b = self.embedding.weight[1:]  # n_nodes x latent_size，309x100
        print('b的大小：', b.shape)  # # n_nodes x latent_size，309x100

        # ======= 断点调试 ======
        # 命令行断点调试，查看u_A_in和u_A
        # pdb.set_trace()
        # ======= 断点调试 ======

        # 原文中公式(8)，b.transpose(1, 0)表示第1维和第0维交换；如果是b.transpose(0, 1)则不变，因为是按照原坐标轴改变序列，内容不变
        scores = torch.matmul(a, b.transpose(1, 0))  # n_nodes x latent_size，309x100
        # ======= 内容查看 =======
        # print('compute_scores中a的内容。。。。。')
        # print(a[:5])
        # print('compute_scores中b的内容。。。。。')
        # print(b[:5])
        # print('compute_scores中scores的内容。。。。。')
        # print(scores[:5])
        # ======= 内容查看 =======
        # batch_size*item_num（数据集中item的总个数），scores的内容是每个会话对应的候选项的预测分数
        return scores

    def forward(self, inputs, A):
        """

        A: 在input的基础上增加一个长度为2*序列顶点数量的维度，储存各顶点的出入度
        hidden：隐藏层，储存顶点向量。hidden在input两个维度基础上增加了第三个维度，长度为hiddenSize。
        gnn：返回经论文中的更新函数更新后的hidden矩阵
        """
        hidden = self.embedding(inputs)
        hidden = self.gnn(A, hidden)
        # 项目的
        return hidden


# 使用GPU加速
def trans_to_cuda(variable):
    if torch.cuda.is_available():
        return variable.cuda()
    else:
        return variable


# 使用CPU计算
def trans_to_cpu(variable):
    if torch.cuda.is_available():
        return variable.cpu()
    else:
        return variable


def forward(model, i, data):
    # 获得数据集中数据第i个batch的详细信息
    # get_slice返回：点击序列对应唯一动作集合的位置角标，该批数据图矩阵的列表，该batch的点击序列矩阵，mask矩阵，目标数据
    alias_inputs, A, items, mask, targets = data.get_slice(i)
    # 将数据放到GPU上
    alias_inputs = trans_to_cuda(torch.Tensor(alias_inputs).long())
    items = trans_to_cuda(torch.Tensor(items).long())
    A = trans_to_cuda(torch.Tensor(A).float())
    mask = trans_to_cuda(torch.Tensor(mask).long())
    # 获得项目的embedding
    hidden = model(items, A)
    print('hidden的大小：', hidden.shape)
    print('alias_inputs的大小：', alias_inputs.shape)
    # ======= 断点测试 =====
    # 查看scores和targets内容
    # pdb.set_trace()
    # ======= 断点测试 =====
    # 定义一个函数，获得对应的item embedding，选择这一批第i个样本对应类别序列的函数，获得train_data中第i个batch的第j个序列的所有项目embedding
    get = lambda j: hidden[j][alias_inputs[j]]
    # 通过stack函数扩维拼接得到会话的embedding
    # torch.stack()同样是对tensors沿指定维度拼接，但返回的Tensor会多一维(形象的理解：假如数据都是二维矩阵(平面)，它可以把这些一个个平面按第三维(例如：时间序列)压成一个三维的立方体，而立方体的长度就是时间序列长度。)
    seq_hidden = torch.stack([get(k) for k in torch.arange(len(alias_inputs)).long()])
    # 给出目标值和计算之后的得分。seq_hidden是一个batch所有序列的隐层输出，batch_size*seqLength*latent_size(hidden_size)
    return targets, model.compute_scores(seq_hidden, mask)


def train_test(model, train_data, test_data):
    # 如果要调整学习率，通常每个epoch更新一次学习率（但不绝对）
    model.scheduler.step()
    # ========= 开始训练 =========
    print('start training: ', datetime.datetime.now())
    model.train()
    total_loss = 0.0
    # 划分batch
    slices = train_data.generate_batch(model.batch_size)
    for i, j in zip(slices, np.arange(len(slices))):
        # 初始化梯度为0
        model.optimizer.zero_grad()
        # 计算预测的目标值和对应的分数，传入模型model(SessionGraph), 数据批的索引i, 训练的数据data(Data);返回目标数据、每个候选项对应的分数
        targets, scores = forward(model, i, train_data)
        targets = trans_to_cuda(torch.Tensor(targets).long())
        # 计算损失
        loss = model.loss_function(scores, targets - 1)
        # 反向传播
        loss.backward()
        # 更新梯度
        model.optimizer.step()
        # 一个epoch整体的损失
        total_loss += loss
        # 每6个batch输出一次损失
        if j % int(len(slices) / 5 + 1) == 0:
            print('[%d/%d] Loss: %.4f' % (j, len(slices), loss.item()))
        # print('[%d/%d] batch Loss: %.4f' % (j, len(slices), loss.item()))
    print('\tCurr Epoch\'s Loss:\t%.3f' % total_loss)
    # ========= 结束训练 =========

    # ======== 开始测试 ==========
    print('start predicting: ', datetime.datetime.now())
    # 不再更新梯度
    model.eval()
    hit, mrr = [], []
    # 划分测试集batch
    slices = test_data.generate_batch(model.batch_size)
    # 计算HR指标和MRR指标
    for i in slices:
        # targets的shape为(batch_size-1,)给定的目标值；scores的shape为batch_size*item_num（数据集中item的总个数），scores的内容是每个会话对应的候选项的预测分数
        targets, scores = forward(model, i, test_data)
        print('scores的shape：', scores.shape)
        # 返回候选项中分数最高的前20个item的索引（其实就是itemId-1），这里索引是item生成embedding时的数组索引，所以索引和itemId差值为1，即index=itemId-1，所以计算hit时需要target-1
        sub_scores = scores.topk(20)[1]
        sub_scores = trans_to_cpu(sub_scores).detach().numpy()

        # ======= 断点测试 =====
        # 查看scores和targets内容
        pdb.set_trace()
        # ======= 断点测试 =====

        for score, target, mask in zip(sub_scores, targets, test_data.mask):
            hit.append(np.isin(target - 1, score))
            if len(np.where(score == target - 1)[0]) == 0:
                mrr.append(0)
            else:
                mrr.append(1 / (np.where(score == target - 1)[0][0] + 1))
    hit = np.mean(hit) * 100
    mrr = np.mean(mrr) * 100
    # ======== 结束测试 ==========
    return hit, mrr
